#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <math.h>
#include <string.h>
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_statistics.h>
#include "cpgplot.h"

#define C(i) (gsl_vector_get(c,(i)))
#define COV(i,j) (gsl_matrix_get(cov,(i),(j)))

#define BUFSIZE 256

void minmax(int nvals,double *minimum,double *maximum,...);

int main(int argc,char *argv[]){
  FILE *input_data=NULL;
  char inp_line[BUFSIZE];
  int nfile,num_files=-1,*n_data_read=NULL,i,j,k,n_all,nfits=100;
  double inp_freq,inp_flux,**frequencies_read=NULL,**fluxes_read=NULL;
  double **rms_calculated=NULL,*all_freqs=NULL,*all_fluxes=NULL,*all_variances=NULL;
  double min_freq,max_freq,min_flux,max_flux;

  /* the fitting variables */
  double chisq;
  gsl_matrix *X=NULL,*cov=NULL;
  gsl_vector *y=NULL,*wt=NULL,*c=NULL;
  gsl_multifit_linear_workspace *work=NULL;
  float *fy_fit_line=NULL,*fx_fit_line=NULL;
  double *residual_flux=NULL;
  double mean,variance,constraint;
  int fit_order;

  /* load the data */
  input_data=fopen("combined_spectra.txt","r");
  while(fgets(inp_line,BUFSIZE,input_data)!=NULL){
    sscanf(inp_line,"%d %lf %lf",&nfile,&inp_freq,&inp_flux);
    if (nfile>num_files){
      num_files=nfile;
      n_data_read=realloc(n_data_read,(num_files+1)*sizeof(int));
      n_data_read[num_files]=0;
      frequencies_read=realloc(frequencies_read,(num_files+1)*sizeof(double *));
      frequencies_read[num_files]=NULL;
      fluxes_read=realloc(fluxes_read,(num_files+1)*sizeof(double *));
      fluxes_read[num_files]=NULL;
      rms_calculated=realloc(rms_calculated,(num_files+1)*sizeof(double *));
      rms_calculated[num_files]=NULL;
    }
    n_data_read[nfile]++;
    frequencies_read[nfile]=realloc(frequencies_read[nfile],
				    n_data_read[nfile]*sizeof(double));
    fluxes_read[nfile]=realloc(fluxes_read[nfile],
			       n_data_read[nfile]*sizeof(double));
    rms_calculated[nfile]=realloc(rms_calculated[nfile],
				  n_data_read[nfile]*sizeof(double));
    frequencies_read[nfile][n_data_read[nfile]-1]=inp_freq;
    fluxes_read[nfile][n_data_read[nfile]-1]=inp_flux;
  }
  fclose(input_data);

  /* work out the RMS of each set by making a straight line fit to the 
     data, then subtracting that fit and working out the residual scatter */
  cpgopen("43/xs");
  cpgask(0);
  for (i=0;i<=num_files;i++){
    X=gsl_matrix_alloc(n_data_read[i],2);
    y=gsl_vector_alloc(n_data_read[i]);
    wt=gsl_vector_alloc(n_data_read[i]);
    c=gsl_vector_alloc(2);
    cov=gsl_matrix_alloc(2,2);
    min_freq=max_freq=frequencies_read[i][0];
    min_flux=max_flux=fluxes_read[i][0];
    for (j=0;j<n_data_read[i];j++){
      gsl_matrix_set(X,j,0,1.0);
      gsl_matrix_set(X,j,1,frequencies_read[i][j]);
      gsl_vector_set(y,j,fluxes_read[i][j]);
      gsl_vector_set(wt,j,1.0);
      minmax(3,&min_freq,&max_freq,min_freq,max_freq,frequencies_read[i][j]);
      minmax(3,&min_flux,&max_flux,min_flux,max_flux,fluxes_read[i][j]);
    }
    work=gsl_multifit_linear_alloc(n_data_read[i],2);
    gsl_multifit_wlinear(X,wt,y,c,cov,&chisq,work);
    gsl_multifit_linear_free(work);
    /* now plot the data and the fit */
    cpgpage();
    cpgsci(1);
    cpgswin(min_freq,max_freq,min_flux,max_flux);
    cpgbox("BCNTS",0.0,0,"BCNTS",0.0,0);
    fx_fit_line=malloc(n_data_read[i]*sizeof(float));
    fy_fit_line=malloc(n_data_read[i]*sizeof(float));
    for (j=0;j<n_data_read[i];j++){
      cpgpt1((float)(frequencies_read[i][j]),(float)(fluxes_read[i][j]),2);
      fx_fit_line[j]=(float)(frequencies_read[i][j]);
      fy_fit_line[j]=(float)(C(0)+C(1)*fx_fit_line[j]);
    }
    cpgsci(2);
    cpgline(n_data_read[i],fx_fit_line,fy_fit_line);

    /* now subtract the fit from each point, and calculate the RMS scatter */
    residual_flux=malloc(n_data_read[i]*sizeof(double));
    for (j=0;j<n_data_read[i];j++){
      residual_flux[j]=(double)(fluxes_read[i][j]-(double)fy_fit_line[j]);
      if (j==0){
	min_flux=max_flux=residual_flux[j];
      } else {
	minmax(3,&min_flux,&max_flux,min_flux,max_flux,residual_flux[j]);
      }
    }
    mean=gsl_stats_mean(residual_flux,(size_t)1,(size_t)n_data_read[i]);
    variance=gsl_stats_variance_m(residual_flux,(size_t)1,(size_t)n_data_read[i],mean);
    printf("mean value = %lg, variance = %lg\n",mean,variance);
    cpgpage();
    cpgsci(1);
    cpgswin(min_freq,max_freq,min_flux,max_flux);
    cpgbox("BCNTS",0.0,0,"BCNTS",0.0,0);
    for (j=0;j<n_data_read[i];j++){
      cpgpt1((float)(frequencies_read[i][j]),(float)(residual_flux[j]),2);
      rms_calculated[i][j]=variance;
    }

    /* free our memory */
    free(fx_fit_line); fx_fit_line=NULL;
    free(fy_fit_line); fy_fit_line=NULL;
    free(residual_flux); residual_flux=NULL;
    gsl_matrix_free(X);
    gsl_vector_free(y);
    gsl_vector_free(wt);
    gsl_vector_free(c);
    gsl_matrix_free(cov);
    
  }

  /* make a giant array of all points */
  n_all=0;
  constraint=1;
  for (i=0;i<=num_files;i++){
    if (rms_calculated[i][0]<constraint){
      printf("file %d, rms = %lg\n",i,rms_calculated[i][0]);
      n_all+=n_data_read[i];
    }
  }
  all_freqs=malloc(n_all*sizeof(double));
  all_fluxes=malloc(n_all*sizeof(double));
  all_variances=malloc(n_all*sizeof(double));
  for (i=0,k=0;i<=num_files;i++){
    if (rms_calculated[i][0]<constraint){
      for (j=0;j<n_data_read[i];j++,k++){
	all_freqs[k]=log10(frequencies_read[i][j]*1000.0);
	all_fluxes[k]=log10(fluxes_read[i][j]);
	all_variances[k]=rms_calculated[i][j];
      }
    }
  }
  /* fit to this: RJS uses a 3rd order fit, so do we */
  fit_order=2;
  X=gsl_matrix_alloc(n_all,fit_order);
  y=gsl_vector_alloc(n_all);
  wt=gsl_vector_alloc(n_all);
  c=gsl_vector_alloc(fit_order);
  cov=gsl_matrix_alloc(fit_order,fit_order);
  min_freq=max_freq=all_freqs[0];
  min_flux=max_flux=all_fluxes[0];
  for (i=0;i<n_all;i++){
    gsl_matrix_set(X,i,0,1.0);
    gsl_matrix_set(X,i,1,all_freqs[i]);
/*     gsl_matrix_set(X,i,2,pow(all_freqs[i],2)); */
/*     gsl_matrix_set(X,i,3,pow(all_freqs[i],3)); */
    gsl_vector_set(y,i,all_fluxes[i]);
/*     gsl_vector_set(wt,i,1.0/pow(all_variances[i],2)); */
/*     if (all_freqs[i]<1.4149){ /\* about 26 GHz *\/ */
    if (all_freqs[i]<4.41497){ /* about 26000 MHz */
      gsl_vector_set(wt,i,10.0/pow(all_variances[i],2)); /* higher weights */
    } else {
      gsl_vector_set(wt,i,1.0/pow(all_variances[i],2));
    }
    minmax(3,&min_freq,&max_freq,min_freq,max_freq,all_freqs[i]);
    minmax(3,&min_flux,&max_flux,min_flux,max_flux,all_fluxes[i]);
  }
  work=gsl_multifit_linear_alloc(n_all,fit_order);
  gsl_multifit_wlinear(X,wt,y,c,cov,&chisq,work);
  gsl_multifit_linear_free(work);

  /* what is the fit? */
/*   printf("best fit: log(flux) = %g + %g log(freq) + %g log(freq)^2 + %g log(freq)^3\n", */
/* 	 C(0),C(1),C(2),C(3)); */
  printf("best fit: log(flux) = %g + %g log(freq)\n",C(0),C(1));

  /* plot everything */
  cpgpage();
  cpgsci(1);
  cpgswin(min_freq,max_freq,min_flux,max_flux);
  cpgbox("BCLNTS",0.0,0,"BCLNTS",0.0,0);
  for (i=0;i<n_all;i++){
    cpgpt1((float)(all_freqs[i]),(float)(all_fluxes[i]),2);
  }    
  fx_fit_line=malloc(nfits*sizeof(float));
  fy_fit_line=malloc(nfits*sizeof(float));
  for (i=0;i<nfits;i++){
    fx_fit_line[i]=min_freq+i*((max_freq-min_freq)/nfits);
/*     fy_fit_line[i]=C(0)+C(1)*fx_fit_line[i]+C(2)*powf(fx_fit_line[i],2)+ */
/*       C(3)*powf(fx_fit_line[i],3); */
    fy_fit_line[i]=C(0)+C(1)*fx_fit_line[i];
  }
  cpgsci(2);
  cpgline(nfits,fx_fit_line,fy_fit_line);

  /* free our memory */
  free(fx_fit_line); fx_fit_line=NULL;
  free(fy_fit_line); fy_fit_line=NULL;
  gsl_matrix_free(X);
  gsl_vector_free(y);
  gsl_vector_free(wt);
  gsl_vector_free(c);
  gsl_matrix_free(cov);

  cpgclos();
  
  exit(0);
}

void minmax(int nvals,double *minimum,double *maximum,...){
  /* returns the minimum and maximum values from a list of values */
  va_list ap;
  int i;
  double nval;
  
  va_start(ap,nvals);
  nval=va_arg(ap,double);
  *minimum=nval;
  *maximum=nval;
  for (i=1;i<nvals;i++){
    nval=va_arg(ap,double);
    if (nval<*minimum){
      *minimum=nval;
    }
    if (nval>*maximum){
      *maximum=nval;
    }
  }
  va_end(ap);

}
